from collections import defaultdict, deque
from typing import Dict, List
import numpy as np
from numpy.core.fromnumeric import mean
from numpy.lib.function_base import select
from pandas.core.frame import DataFrame
from src.datahelper import StockDataHelper
from src.backTestingData import BackTestingData
from src.taskSignals import TaskSignals
from src.strategy.baseStrategy import BaseStrategy


class DoubleMaStrategy(BaseStrategy):
    def __init__(self) -> None:
        super().__init__()

    def init_strategy(self, dataHelper: StockDataHelper,
                      backTestingData: BackTestingData, taskSignals: TaskSignals):
        self.taskSignals = taskSignals
        self.backTestingData = backTestingData
        self.data_by_date = dataHelper.get_start_end_dayData(
            backTestingData.data['start_date'], backTestingData.data['end_date'])
        self.trade_day_list = list(self.data_by_date.keys())
        self.service_charge = backTestingData.data['service_charge']
        self.position.clear()
        self.daily_return.clear()
        self.money = self.backTestingData.data['money']
        self.initial_money = self.money
        # 初始买入的股票
        self.buy_stocks(
            backTestingData.data['selected_stocks'], self.trade_day_list[0], 'open')
        self.update_portfolio_value(self.trade_day_list[0])

    def run(self):
        long_ma = self.backTestingData.data['long_ma']
        short_ma = self.backTestingData.data['short_ma']

        # 维护长短均线内所有的日期窗口
        long_ma_window = defaultdict(deque)
        short_ma_window = defaultdict(deque)
        # 计算长短窗口内的所有股票的价值和
        long_ma_sum = defaultdict(float)
        short_ma_sum = defaultdict(float)
        # 上一天的ma
        last_ma = None
        last_tot = 0
        for i, date in enumerate(self.trade_day_list):
            stocks = list(self.data_by_date[date].index)
            long_ma_start_date = StockDataHelper.get_pre_day(date, long_ma)
            short_ma_start_date = StockDataHelper.get_pre_day(date, short_ma)
            for stock in stocks:
                # 当天的收盘价
                price = self.data_by_date[date].loc[stock, 'close']
                if np.isnan(price):
                    continue
                long_ma_window[stock].append(date)
                long_ma_sum[stock] += price
                short_ma_window[stock].append(date)
                short_ma_sum[stock] += price

                # 弹出窗口之外的日期
                while len(long_ma_window[stock]) > 0 and long_ma_window[stock][0] < long_ma_start_date:
                    long_ma_window[stock].popleft()
                    long_ma_sum[stock] -= price

                # 弹出窗口之外的日期
                while len(short_ma_window[stock]) > 0 and short_ma_window[stock][0] < short_ma_start_date:
                    short_ma_window[stock].popleft()
                    short_ma_sum[stock] -= price

            ma_data = {
                'date': date,
                'short_ma': defaultdict(float),
                'long_ma': defaultdict(float)
            }
            # 均值
            for stock in long_ma_sum.keys():
                ma_data['long_ma'][stock] = long_ma_sum[stock] / \
                    len(long_ma_window[stock])
                ma_data['short_ma'][stock] = short_ma_sum[stock] / \
                    len(short_ma_window[stock])

            eps = 1e-8
            if last_ma != None:
                buys_stocks, sell_stocks = [], []
                for stock in stocks:
                    price = self.data_by_date[date].loc[stock, 'close']
                    if np.isnan(price):
                        continue
                    if np.fabs(ma_data['long_ma'][stock]) < eps or np.fabs(ma_data['short_ma'][stock]) < eps:
                        continue

                    if last_ma['short_ma'][stock] < last_ma['long_ma'][stock] and ma_data['short_ma'][stock] > ma_data['long_ma'][stock]:
                        buys_stocks.append(stock)

                    if last_ma['short_ma'][stock] > last_ma['long_ma'][stock] and ma_data['short_ma'][stock] < ma_data['long_ma'][stock]:
                        sell_stocks.append(stock)

                self.buy_stocks(buys_stocks, date)
                self.sell_stocks(sell_stocks, date)

            last_ma = ma_data
            self.update_portfolio_value(date)
            self.taskSignals.progressBarSignal.emit(
                int(i / len(self.trade_day_list) * 100))

            cur = self.money + self.portfolio_value
            if last_tot != 0:
                self.daily_return.append(
                    (cur - last_tot) / self.initial_money * 100)
            last_tot = cur

        self.taskSignals.progressBarSignal.emit(100)

        self.log(f'最终资金{self.money + self.portfolio_value}')
        self.log(f'最大回撤{StockDataHelper.max_drawdown(self.daily_return)}')

        self.updateBackTestingData()
